from typing import List, Dict
from collections import OrderedDict
from datetime import datetime, timedelta
from src.api_calls import query_paper_metadata_that_title_contain, query_by_paper_id
import re

class PaperCache:
    def __init__(self, max_size=1000):
        self.cache = OrderedDict()
        self.max_size = max_size

    def get(self, paper_id):
        if paper_id in self.cache:
            self.cache.move_to_end(paper_id)
            return self.cache[paper_id]
        return None

    def set(self, paper_id, data):
        if paper_id not in self.cache and len(self.cache) >= self.max_size:
            self.cache.popitem(last=False)
        self.cache[paper_id] = data

class RetrievalAgent:
    def __init__(self):
        self.paper_cache = PaperCache()
        self.citation_counter = 1
        self.citation_map = {}
        # 会议名称映射表
        self.venue_map = {
            "conf": "Conference",
            "nips": "NeurIPS",
            "neurips": "NeurIPS",
            "icml": "ICML",
            "iclr": "ICLR",
            "aaai": "AAAI",
            "cvpr": "CVPR",
            "acl": "ACL",
            "emnlp": "EMNLP",
            "naacl": "NAACL",
            "eccv": "ECCV",
            "ijcai": "IJCAI",
            "kdd": "KDD",
            "www": "WWW",
            "sigir": "SIGIR",
            "arxiv": "arXiv preprint",
            "preprint": "Preprint",
            "wsdm": "WSDM",
            "uai": "UAI",
            "aistats": "AISTATS",
            "colt": "COLT",
            "usenix": "USENIX",
            "sigmod": "SIGMOD",
            "vldb": "VLDB",
            "icde": "ICDE",
            "coling": "COLING",
            "interspeech": "INTERSPEECH",
            "mm": "ACM MM",
            "iccv": "ICCV"
        }
    def _normalize_venue(self, raw_venue: str) -> str:
        """标准化会议名称处理"""
        if not raw_venue or raw_venue.isspace():
            return "Conference"
        venue_lower = raw_venue.lower()
        return self.venue_map.get(venue_lower, raw_venue)
        
    def _match_arxiv_to_conference(self, arxiv_id: str) -> str:
        """将arXiv ID映射到实际会议"""
        conf_info = self.arxiv_conf_map.get(arxiv_id)
        if conf_info:
            conf_name, year = conf_info
            return self.venue_map.get(conf_name.lower(), "arXiv preprint")
        return "arXiv preprint"
    def _decompose_query(self, topic: str) -> List[str]:
        """基于规则的关键词提取，将用户输入的主题分解为关键词组合"""
        try:
            # 移除常见的停用词和标点符号
            stop_words = {'and', 'or', 'the', 'in', 'on', 'at', 'to', 'for', 'of', 'with'}
            cleaned_topic = ''.join(c.lower() if c.isalnum() else ' ' for c in topic)
            
            # 分词并过滤
            words = [word.strip() for word in cleaned_topic.split()]
            keywords = [word for word in words 
                       if word and len(word) > 2 
                       and word not in stop_words]
            
            # 去重并限制关键词数量
            unique_keywords = list(dict.fromkeys(keywords))
            return unique_keywords[:5]
            
        except Exception as e:
            print(f"关键词提取失败: {str(e)}")
            return [topic]

    def search_papers(self, query: str, top_k: int = 30) -> List[Dict]:
        """通过papers API搜索并确保返回结果包含必要字段"""
        try:
            from src.api_calls import search_papers
            papers = search_papers(query, top_k=top_k)
            
            results = []
            for paper in papers:
                if not paper.get('entity'):
                    print(f"跳过缺失entity的论文")
                    continue
                    
                entity = paper['entity']
                paper_id = entity.get('paper_id')
                title = entity.get('paper_title')
                chunk_id = entity.get('chunk_id')
                # 从original_filename和其他字段中提取会议信息
                original_filename = entity.get('original_filename', '')
                raw_venue = entity.get('venue', '')
                venue = 'Conference'
                
                # 首先尝试从venue字段提取
                if raw_venue:
                    venue = self._normalize_venue(raw_venue)
                
                # 如果venue仍为Conference，尝试从文件名提取
                if venue == 'Conference' and original_filename:
                    # 改进的会议名称匹配模式
                    conf_patterns = [
                        r'([A-Za-z]+)(?:20\d{2})',  # 标准格式 CONF2023
                        r'(?:^|[^a-zA-Z])([A-Za-z]+)(?:[^a-zA-Z]|$)',  # 独立的会议缩写
                        r'([A-Za-z]+)[-_]\d{2,4}'  # 带分隔符的格式
                    ]
                    
                    for pattern in conf_patterns:
                        match = re.search(pattern, original_filename, re.IGNORECASE)
                        if match:
                            conf_name = match.group(1)
                            normalized_venue = self._normalize_venue(conf_name)
                            if normalized_venue != 'Conference':
                                venue = normalized_venue
                                break
                
                print(f"处理论文: {title}")
                
                # 优化数据验证逻辑，只要title存在就保留
                if not title:
                    print(f"跳过缺失标题的论文: paper_id={paper_id}")
                    continue
                    
                # 为缺失字段设置默认值
                if not paper_id:
                    paper_id = f"unknown_{len(results)}"
                if not chunk_id:
                    chunk_id = str(len(results))
                    
                unique_id = f"{paper_id}_{chunk_id}"
                if unique_id not in self.citation_map:
                    self.citation_map[unique_id] = self.citation_counter
                    self.citation_counter += 1
                
                results.append({
                    'title': title,
                    'paper_id': paper_id,
                    'chunk_id': chunk_id,
                    'content': entity.get('chunk_text', ''),
                    'venue': venue,
                    'year': str(entity.get('year', datetime.now().year)),
                    'unique_id': unique_id,
                    'citation_number': self.citation_map[unique_id],
                    'entity': entity  # 添加完整的entity信息以供后续使用
                })
                
                if len(results) >= top_k:
                    break
            
            return results
            
        except Exception as e:
            print(f"论文检索失败: {str(e)}")
            return []

    def get_citation_number(self, unique_id: str) -> int:
        """获取引用编号"""
        return self.citation_map.get(unique_id)

    def format_citation(self, unique_id: str) -> str:
        """格式化引用标记"""
        citation_number = self.get_citation_number(unique_id)
        if citation_number:
            return f"<sup>{citation_number}</sup>"
        return ""
